from collections import deque
from dataclasses import dataclass
from typing import Union
import logging
from coin.base.timestamp import get_timestamp

INITIAL_VOLATILITY = 10


@dataclass
class Entry:
    entry_num: int
    plain_sum: float
    power2_sum: float
    timestamp_start: Union[None, int] = None

    def add(self: "Entry", other: "Entry") -> None:
        self.entry_num += other.entry_num
        self.plain_sum += other.plain_sum
        self.power2_sum += other.power2_sum

    def minus(self: "Entry", other: "Entry") -> None:
        self.entry_num -= other.entry_num
        self.plain_sum -= other.plain_sum
        self.power2_sum -= other.power2_sum

    def copy_from(self: "Entry", other: "Entry") -> None:
        self.entry_num = other.entry_num
        self.plain_sum = other.plain_sum
        self.power2_sum = other.power2_sum
        self.timestamp_start = other.timestamp_start

    def volatilty(self: "Entry") -> float:
        if self.entry_num > 0:
            mean = 1.0 * self.plain_sum / self.entry_num
            p2mean = 1.0 * self.power2_sum / self.entry_num
            retval = (p2mean - (mean**2)) / (self.entry_num) # take as given the interval of entries are fixed
            
            retval = 0 if abs(retval) < 1e-4 else retval
            if (retval < 0):
                retval = INITIAL_VOLATILITY
            return retval
        return None


class volatility_estimator(object):
    SHORT_TERM_SIZE = 15  # 15 minutes
    SLOT_TIMEINTERVAL = 10  # seconds
    LONG_TERM_SIZE = 240  # 4 hours

    def __init__(self):
        self.hypothetical_volatility_estimate = None
        self.data_queue = deque([Entry(0, 0, 0, None) for i in range(self.LONG_TERM_SIZE)])
        self.short_term = Entry(0, 0, 0, None)
        self.long_term = Entry(0, 0, 0, None)
        self.volatility_bandwidth = 0.25  # stay relatively consistent with simple amm main strategy logic
        
        self.prevPrice = None

    @property
    def short_term_volatility(self):
        return self.short_term.volatilty()

    @property
    def long_term_volatility(self):
        return self.long_term.volatilty()

    def need_new_blk(self, timestamp: int):
        current_timestamp = self.data_queue[-1].timestamp_start
        return (current_timestamp is not None) and (current_timestamp + self.SLOT_TIMEINTERVAL * (10 ** 9) < timestamp)

    def update_volatility(self, price_input: float, timestamp: int) -> None:
        if self.prevPrice is None:
            self.prevPrice = price_input

        _tmp = price_input
        price_input = 1.0e6 * (1.0 * price_input / self.prevPrice) - 1
        self.prevPrice = _tmp
        
        current_entry = Entry(1, price_input, price_input**2, timestamp)

        if (self.need_new_blk(timestamp)):
            # take a new blk from left
            new_blk = self.data_queue.popleft()
            self.long_term.minus(new_blk)

            # remove short term blk from left
            left_st = self.data_queue[-self.SHORT_TERM_SIZE]
            self.short_term.minus(left_st)

            # append new blk to the right
            new_blk.copy_from(current_entry)
            self.long_term.add(current_entry)
            self.short_term.add(current_entry)
            self.data_queue.append(current_entry)
        else:
            if self.data_queue[-1].timestamp_start is None:
                self.data_queue[-1].timestamp_start = current_entry.timestamp_start
            self.data_queue[-1].add(current_entry)
            self.short_term.add(current_entry)
            self.long_term.add(current_entry)

    def hypothetical_volatility_estimate_correct(self) -> Union[None, bool]:
        if self.hypothetical_volatility_estimate is None:
            return None

        short_term_vol = self.short_term.volatilty()

        return \
          short_term_vol > self.hypothetical_volatility_estimate * (1 - self.volatility_bandwidth) and \
          short_term_vol < self.hypothetical_volatility_estimate * (1 + self.volatility_bandwidth)

    def set_assumed_volatility_estimate(self, hypothesis: float):
        self.hypothetical_volatility_estimate = hypothesis

    def on_book(self, price: float, timestamp: int):
        self.update_volatility(price, timestamp)
        logging.info(f"sv:{self.short_term_volatility}, lv:{self.long_term_volatility}")
